# -*- coding: utf-8 -*-
"""

PBAR_CheckBasicScanZspecFeatureCalculation.py

Created on Tue May 06 17:35:38 2014

@author: jkwong
"""

import PBAR_Zspec, PBAR_Cargo
reload(PBAR_Zspec)
import numpy as np
import os,cPickle,copy, glob
import matplotlib.pyplot as plt
#from scipy import interpolate
from matplotlib import cm
from scipy import ndimage


# Set useful plot variables
plotColors = ['r', 'b', 'g', 'm', 'c', 'y', 'k'] * 10
lineStyles = ['-', '-.', ':', '_', '|'] *  10
markerTypes = ['.', 'o', 'v', '^', '<', '>', '1', '2', '3', '4', 's', 'p', '*', 'h', 'H', '+', 'x', 'D', 'd']
markerTypes = markerTypes * 2
colorList = np.array(['r', 'b', 'g', 'm', 'c', 'y'])

(goodZspecMask, badZspecMask, goodZspecNameList, badZspecNameList, \
        goodZspecIndices, badZspecIndices) = PBAR_Zspec.ZspecDetectorLists2ndSet()

dataPath = r'E:\PBAR\data4\BasicScansStandardWidth'
data0Path = r'E:\PBAR\data4\BasicScanCargo'

plotDir = r'E:\PBAR\data4\BasicScansPlots'
plot0Dir = r'E:\PBAR\data4\BasicScansPlots'

basepathSet2 = r'C:\Users\jkwong\Documents\Work\PBAR\data4\Mar-files'
plotSaveDir = r'E:\PBAR\data4\BasicScansPlots'

datasetDescription = PBAR_Zspec.ReadCargoDataDescriptionFile(r'E:\PBAR\data4\CargoSet2.txt')


#        cargoCountRange = np.array([0.01e8, 0.4e8]) # only consider parts of the zspec image that have rad below this value
cargoCountRange = np.array([0, 0.4e8]) # only consider parts of the zspec image that have rad below this value


acqTime = 1/60.
figureSize = (16, 10)

savePlots = 1
closePlots = 1


# set 2 - load from file
setNum = 2
fullFilename = os.path.join(basepathSet2, 'zspec%dSet.dat' %setNum)
with open(fullFilename ,'rb') as fid:
    print('Reading %s' %fullFilename)
    temp = cPickle.load(fid)

polyMeanOrderExtendIndex = 22
polySigmaOrderExtendIndex = 13

pfitmean = temp['pfit']['mean']['PbALL']['extend'][polyMeanOrderExtendIndex]
pfitsigma = temp['pfit']['sigma']['PbALL']['extend'][polySigmaOrderExtendIndex]


#################################
## DEFINE FILE NAMES

energy = np.array([7.50E+00,1.05E+01,1.35E+01,1.65E+01,1.95E+01,2.25E+01,2.55E+01, \
    2.85E+01,3.15E+01,3.45E+01,3.75E+01,4.05E+01,4.35E+01,4.65E+01,4.95E+01,\
    5.25E+01,5.55E+01,5.85E+01,6.15E+01,6.45E+01,6.75E+01,7.05E+01,7.35E+01, \
    7.65E+01,7.95E+01,8.25E+01,8.55E+01,8.85E+01,9.15E+01,9.45E+01,9.75E+01,1.01E+02])


datasetIndex = 7
print(datasetIndex)

######################
## READ IN DATA

# Read Zspec
#    filenameZspec = '%s-FDFC-All.npy' %datasetDescription['scanID'][datasetIndex]
#    fullFilenameZspec = os.path.join(data0Path, filenameZspec)
#    print('Loading %s' %fullFilenameZspec)
#    datZspec = np.load(fullFilenameZspec)

filenameZspec = '%s-FDFC-All_SW.npy' %datasetDescription['scanID'][datasetIndex]
fullFilenameZspec = os.path.join(dataPath, filenameZspec)
print('Loading %s' %fullFilenameZspec)
datStandardWidth = np.load(fullFilenameZspec)

# Read Cargo
#    filenameCargo = 'PBAR-%s.cargoimage' %datasetDescription['dataFile'][datasetIndex]
#    fullFilenameCargo = os.path.join(dataPath, filenameCargo)
#    print('Loading %s' %fullFilenameCargo)
##    datCargo= np.load(fullFilenameCargo)
#    (datCargo,bpp,formatt,flag,low1,high1,low2,high2) = PBAR_Cargo.ReadCargoImage(fullFilenameCargo)

filenameCargo = 'PBAR-%s.cargoimageSW.npy' %datasetDescription['dataFile'][datasetIndex]
fullFilenameCargo = os.path.join(dataPath, filenameCargo)
print('Loading %s' %fullFilenameCargo)
datCargoStandardWidth= np.load(fullFilenameCargo)

# Read in marker files
filenameMarker = filenameCargo.replace('cargoimageSW', 'cargomarkerSW')
fullFilenameMarker = fullFilenameCargo.replace('cargoimageSW', 'cargomarkerSW')
# some don't have marker files
if os.path.exists(fullFilenameMarker):
#        markers = PBAR_Cargo.ReadCargoMarker(fullFilenameMarker)
    markerStandardWidth = PBAR_Cargo.ReadCargoMarker(fullFilenameMarker)
else:
    markerStandardWidth = []

#############################
##  CALCULATE FEATURES

# interp version
discrim = PBAR_Zspec.CalculateFeaturesZspecBasic(datStandardWidth, energy, 1)

# calculate for a pixel
x = 448
y = 84
discrimSingle = PBAR_Zspec.CalculateFeaturesZspecSingleSpectrum(datStandardWidth[x, y,:], energy, 1)


print('Comparing values calculated by CalculateFeaturesZspecBasic and CalculateFeaturesZspecSingleSpectrum')
for key in discrimSingle.keys():
    print('%s, %3.3f, %3.3f' %(key, discrimSingle[key], discrim[key][x, y]))


# PLOTS

# zspec counts
plt.figure()
plt.grid()

plt.imshow(np.log(discrim['count'].T), interpolation = 'nearest', aspect='auto', cmap = cm.Greys_r)

for i, mark in enumerate(markerStandardWidth):
    x = mark['rec_left'], mark['rec_right'], mark['rec_right'], mark['rec_left'], mark['rec_left']
    y = mark['rec_bottom'], mark['rec_bottom'], mark['rec_top'], mark['rec_top'], mark['rec_bottom']
    plt.plot(x, y, 'r')
    plt.plot(mark['x'], mark['y'], 'xr')
    if 'left' in mark:
        plt.plot(mark['left']['x'], mark['left']['y'], 'gx')
    if 'right' in mark:
        plt.plot(mark['right']['x'], mark['right']['y'], 'go')
    plt.text(mark['x'], mark['y'],  '%d, %s' %(i, mark['target']), fontsize = 12, color = 'r')

# radiography
plt.figure()
plt.grid()

plt.imshow(np.log(datCargoStandardWidth.T), interpolation = 'nearest', aspect='auto', cmap = cm.Greys_r)

for i, mark in enumerate(markerStandardWidth):
    x = mark['rec_left'], mark['rec_right'], mark['rec_right'], mark['rec_left'], mark['rec_left']
    y = mark['rec_bottom'], mark['rec_bottom'], mark['rec_top'], mark['rec_top'], mark['rec_bottom']
    plt.plot(x, y, 'r')
    plt.plot(mark['x'], mark['y'], 'xr')
    if 'left' in mark:
        plt.plot(mark['left']['x'], mark['left']['y'], 'gx')
    if 'right' in mark:
        plt.plot(mark['right']['x'], mark['right']['y'], 'go')
    plt.text(mark['x'], mark['y'],  '%d, %s' %(i, mark['target']), fontsize = 12, color = 'r')



# zspec counts
plt.figure()
plt.grid()

plt.imshow(np.log(discrim['binSkew'].T), interpolation = 'nearest', aspect='auto', cmap = cm.Greys_r)

for i, mark in enumerate(markerStandardWidth):
    x = mark['rec_left'], mark['rec_right'], mark['rec_right'], mark['rec_left'], mark['rec_left']
    y = mark['rec_bottom'], mark['rec_bottom'], mark['rec_top'], mark['rec_top'], mark['rec_bottom']
    plt.plot(x, y, 'r')
    plt.plot(mark['x'], mark['y'], 'xr')
    if 'left' in mark:
        plt.plot(mark['left']['x'], mark['left']['y'], 'gx')
    if 'right' in mark:
        plt.plot(mark['right']['x'], mark['right']['y'], 'go')
    plt.text(mark['x'], mark['y'],  '%d, %s' %(i, mark['target']), fontsize = 12, color = 'r')



# plot across a certain time slice
plt.figure()
plt.grid()
plt.plot(discrim['binSkew'][448,:])

# compare

# plot several spectra

plt.figure()
plt.grid()
# plot eleven pixels

xBinRange = [320, 331]
yBin = 49
plt.plot(energy, datStandardWidth[xBinRange[0]:xBinRange[1], yBin,:].T)
plt.plot(energy, datStandardWidth[xBinRange[0]:xBinRange[1], yBin,:].mean(0), linewidth = 2)
plt.yscale('log')
plt.xlabel('Bin')
plt.ylabel('Count')
discrimSingle = PBAR_Zspec.CalculateFeaturesZspecSingleSpectrum(datStandardWidth[xBinRange[0]:xBinRange[1], yBin,:].mean(0), energy, 1)

discrim['binSkew'][xBinRange[0]:xBinRange[1],yBin]

discrimSingle = PBAR_Zspec.CalculateFeaturesZspecSingleSpectrum(datStandardWidth[xBinRange[0]:xBinRange[1], yBin,:].mean(0), energy, 1)

print('Comparing values calculated by CalculateFeaturesZspecBasic and CalculateFeaturesZspecSingleSpectrum')
for key in discrimSingle.keys():
    tempMean = discrim[key][xBinRange[0]:xBinRange[1], yBin].mean(0)
    tempStd = discrim[key][xBinRange[0]:xBinRange[1], yBin].std(0)
    
    print('%s, %3.3f, %3.3f, %3.3f, %3.3f%%' %(key, discrimSingle[key], tempMean, tempStd, tempStd/tempMean *100   ))
    print(discrim[key][xBinRange[0]:xBinRange[1], yBin])
    print(' ')



# test reshaping of the datStandardWidth

temp1 = datStandardWidth[:,100:105,:][200:210,:,:]  # 10x5x32
temp2 = datStandardWidth[:,100:105,:][200:210,:,:].reshape((50, 32)) # 50 x 32



